package uwmadison.soot.analysis;

/* Nate Deisinger (ndeisinger@wisc.edu)
 *
 * This file is part of soot-csi, a project to apply lightweight
 * instrumentation to Java programs to facilitate debugging through
 * the Soot analysis framework.
 *
 * Copyright (c) 2014, Nate Deisinger, University of Wisconsin-Madison
 *   Madison, Wisconsin 53715
 *   All Rights Reserved.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc.,
 *   51 Franklin Street, Fifth Floor,
 *   Boston, MA  02110-1301
 *   USA
 * or retrieve version 2.1 at their website:
 *   http://www.gnu.org/licenses/lgpl-2.1.html
 */

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.HashMap;
import java.util.NoSuchElementException;
import java.util.Queue;
import soot.*;
import soot.jimple.*;
import soot.tagkit.SourceLnPosTag;
import soot.tagkit.Tag;
import soot.toolkits.graph.Block;
import soot.toolkits.graph.ZonedBlockGraph;
import soot.util.Chain;
import soot.util.HashChain;
import uwmadison.soot.SootCSI;

/**
 * Soot whole-program transformation to add path tracing instrumentation
 * to classes.  
 */
public class PathAnalysis {

    private PathTrace dummy = new PathTrace(); //To ensure proper builds
    static enum DUMMY_TYPE{ENTRY, EXIT, NON_DUMMY};
    static final int ENTRY_DUMMY = -1;
    static final int EXIT_DUMMY = -2;

    protected static int tempRefs = 0;
    protected static SootClass pathClass, throwableClass;
    protected static SootMethod blankConstructor, 
    nameConstructor, getRef, dumpAllTraces, freeStack;
    protected static SootField pathVal, inUse;

    /**
     * Performs initialization to link Soot objects to our instrumentation code.
     */
    static
    {
        Scene.v().loadClassAndSupport("uwmadison.soot.analysis.PathTrace");
        pathClass = Scene.v().getSootClass("uwmadison.soot.analysis.PathTrace");
        
        Scene.v().loadClassAndSupport("java.lang.Throwable");
        throwableClass = Scene.v().getSootClass("java.lang.Throwable");
        
        nameConstructor = pathClass.getMethod("void <init>(java.lang.String)");
        freeStack = pathClass.getMethodByName("freeStack");
        getRef = pathClass.getMethodByName("getRef");
        dumpAllTraces = pathClass.getMethodByName("dumpAllTraces");
        
        pathVal = pathClass.getField("long recentPath");
        inUse = pathClass.getField("boolean inUse");
    }

    /*for all node n in reverse topological order do
  if n is a leaf then
     NumPaths(n) ← 1
  else
     NumPaths(n) ← 0
     for all edge e of the form n → m do
        Val(e) ← NumPaths(n)
        NumPaths(n) ← NumPaths(n) + NumPaths(m)
     end for
  end if
end for*/

    /**
     * Maps nodes to instrumentation values.  Used in algorithm.
     */
    private HashMap<CFGNode, Integer> instruments;
    private CFGNode enterNode;
    private CFGNode exitNode;

    /**
     * Performs the first step of calculating Ball-Larus weights.
     */
    private int BallLarus_stub(CFGNode node)
    {
        if(instruments.containsKey(node)) return instruments.get(node);
        if (node.getSuccs().size() == 0)
        {
            instruments.put(node, 1);
            return 1;
        }
        else
        {
            instruments.put(node, 0);
            for (CFGNode succ : node.getSuccs())
            {
                CFGEdge ourEdge = node.getSuccEdges().get(succ);
                ourEdge.setVal(instruments.get(node));
                int newNVal = instruments.get(node) + BallLarus_stub(succ);
                instruments.put(node, newNVal);
            }
            return instruments.get(node);
        }
    }

    /**
     * Keeps track of how many blocks trap to a certain unit.
     */
    HashMap<Unit, Integer> trapMap = new HashMap<Unit, Integer>();

    /**
     * Splits traps into per-block ones and adds appropriate CFG edges
     * for exceptional control flow.
     */
    private void setExceptions()
    {
        //Iterate through traps
        Chain<Trap> traps = methodBody.getTraps();
        LinkedList<Trap> trapsToRemove = new LinkedList<Trap>();
        LinkedList<Trap> trapsToAdd = new LinkedList<Trap>();
        for (Trap t : traps)
        {
            if (doDebug) G.v().out.println("Checking trap " + t);
            Collection<CFGNode> nodes = blockNodes.values();
            //Get catch block
            Unit catchHead = t.getHandlerUnit();
            CFGNode catchNode = null;
            for (CFGNode n : nodes)
            {
                if (n.isDummy()) continue;
                if (n.getBlock().getHead() == catchHead)
                {
                    //Got our block!
                    catchNode = n;
                    break;
                }
            }
            if (catchNode == null) continue;//throw new RuntimeException("Couldn't find a catch node for trap " + t);

            //For each trap get list of blocks it traps and split
            boolean foundNode = false;
            Integer trappedNodes = trapMap.get(t.getHandlerUnit());
            if (trappedNodes == null) trappedNodes = 0;
            for (CFGNode n : nodes)
            {
                if (doDebug) G.v().out.println("Checking for trap from node " + n.getBlockID() + " to " + catchNode.getBlockID());
                if (n.isDummy()) continue;
                if(isInTrap(n.getBlock(), t))
                {
                    foundNode = true;
                    //Add edge from block to catch
                    CFGEdge edge = makeCFGEdge(n, catchNode);
                    if (doDebug) G.v().out.println("Made edge " + edge);

                    Trap newTrap = Jimple.v().newTrap(t.getException(), n.getBlock().getHead(), 
                            n.getBlock().getBody().getUnits().getSuccOf(n.getBlock().getTail()), t.getHandlerUnit());
                    trapsToAdd.add(newTrap);
                    //Need to associate trap info as well
                    edge.setTrap(newTrap);
                    trappedNodes++;
                }
            }
            if (doDebug) G.v().out.println("Trap unit: " + t.getHandlerUnit() + ", nodes: " + trappedNodes);
            trapMap.put(t.getHandlerUnit(), trappedNodes);
            if (foundNode) trapsToRemove.add(t);
        }
        traps.addAll(trapsToAdd);
        traps.removeAll(trapsToRemove);
    }

    /**
     * Returns true if a block is within a given trap.
     * Probably not very efficient.
     * @return Whether or not the block is in the given trap.
     */
    private boolean isInTrap(Block b, Trap t)
    {
        Unit startUnit = t.getBeginUnit();
        Unit endUnit = t.getEndUnit();
        Unit currUnit = startUnit;
        Unit blockHead = b.getHead();
        Body methodBody = b.getBody();
        while (currUnit != endUnit)
        {
            if (currUnit == blockHead) return true;
            currUnit = methodBody.getUnits().getSuccOf(currUnit);
        }
        return false;
    }

    /**
     * Stores our backedges for later use.
     */
    private HashSet<CFGEdge> backEdges;

    /**
     * Removes backedges to create a DAG and adds dummy edges as needed.
     */
    private void fixBackedges()
    {
        //Do a DFS from the start node and find backedges
        backEdges = new HashSet<CFGEdge>();
        backedgeDFS(new HashSet<CFGNode>(), backEdges, new HashSet<CFGNode>(), enterNode);

        //Per BL '96, for backedge w->v, add dummy edges ENTRY->v and w->EXIT and remove w->v
        for (CFGEdge back : backEdges)
        {
            if (doDebug) G.v().out.println("Removing backedge " + back);
            CFGEdge inDummy = makeCFGEdge(enterNode, back.getTo());
            CFGEdge outDummy = makeCFGEdge(back.getFrom(), exitNode);
            inDummy.setDummyEdge(true);
            outDummy.setDummyEdge(true);
            removeCFGEdge(back);
            back.getFrom().addBackedge(back);
            back.setBackedge(true);
            back.setAssociatedDummyEdges(inDummy, outDummy);
        }
    }

    /**
     * Special DFS for finding back edges.
     */
    private void backedgeDFS(HashSet<CFGNode> visited, HashSet<CFGEdge> backEdges, 
            HashSet<CFGNode> seen, CFGNode curr)
    {
        if (visited.contains(curr)) return;
        visited.add(curr);
        seen.add(curr);
        for (CFGEdge succ : curr.getSuccEdges().values())
        {
            if (seen.contains(succ.getTo()))
            {
                //Found a backedge
                backEdges.add(succ);
            }
            else
            {
                backedgeDFS(visited, backEdges, seen, succ.getTo());
            }
        }
        seen.remove(curr);
    }

    /**
     * Removes an edge from the CFG.
     */
    private void removeCFGEdge(CFGEdge e)
    {
        e.getFrom().getSuccEdges().remove(e);
        e.getFrom().getSuccs().remove(e.getTo());
        e.getTo().getPredEdges().remove(e);
        e.getTo().getPreds().remove(e.getTo());
        edges.remove(e);
    }

    /**
     * Set of all non-backedges in graph.
     */
    private HashSet<CFGEdge> edges = new HashSet<CFGEdge>();

    /**
     * Calculates final instrumentation values for chords of a spanning tree.
     */
    private void calculateInstruments(HashSet<CFGEdge> spanningTree)
    {
        //For each chord, calculate instrumentation.
        //Per BL '94: this is sum of values on loop of chord, w/r/t  direction
        HashSet<CFGEdge> chords = new HashSet<CFGEdge>();
        chords.addAll(edges);
        chords.removeAll(spanningTree);
        for (CFGEdge e : chords)
        {
            e.setInstrVal(0);
        }
        chordCycleDFS(0, enterNode, null, spanningTree, chords);
        for (CFGEdge e : chords)
        {
            if (e.isBackedge()) throw new RuntimeException("Error: backedge in chords!");
            e.setInstrVal(e.getInstrVal() + e.getVal());
            if (doDebug) G.v().out.println("Chord " + e + ", final value: " + e.getInstrVal());
            
        }
    }

    /**
     * Special DFS for finding chord cycles.
     * Taken from BL '94.
     */
    private void chordCycleDFS(int initVal, CFGNode node,
            CFGEdge edge, HashSet<CFGEdge> spanningTree, HashSet<CFGEdge> chords)
    {
        if (doDebug) G.v().out.println("InitVal " + initVal + ", node: " + node);
        for (CFGEdge in : node.getPredEdges().values())
        {
            if (spanningTree.contains(in) && (in != edge))
            {
                chordCycleDFS(direction(edge, in)*initVal + in.getVal(), in.getFrom(), in, spanningTree, chords);
            }
        }
        for (CFGEdge out : node.getSuccEdges().values())
        {
            if (spanningTree.contains(out) && (out != edge))
            {
                chordCycleDFS(direction(edge, out)*initVal + out.getVal(), out.getTo(), out, spanningTree, chords);
            }
        }
        for (CFGEdge chord : chords)
        {
            if ((chord.getFrom() == node) || (chord.getTo() == node))
            {
                chord.setInstrVal(chord.getInstrVal() + direction(edge, chord)*initVal);
                if (doDebug) G.v().out.println("InitVal: " + initVal + ", Chord: " + chord + ", instrVal " + chord.getInstrVal());
            }
        }
    }

    /**
     * Returns if two edges sharing an endpoint are in the same direction (1) or not (-1).
     */
    private int direction(CFGEdge a, CFGEdge b)
    {
        if (a == null) return 1;
        if ((a.getTo() == b.getFrom()) || (a.getFrom() == b.getTo())) return 1;
        return -1;
    }

    /**
     * Builds tree of edges that contain increment weights.
     */
    public static HashSet<CFGEdge> buildWeightTree(ArrayList<CFGEdge> edges)
    {
        HashSet<CFGEdge> treeEdges = new HashSet<CFGEdge>();
        for (CFGEdge e : edges)
        {
            if (e.getVal() != 0)
            {
                treeEdges.add(e);
            }
        }
        return treeEdges;
    }

    HashSet<Trap> visitedTraps = new HashSet<Trap>();
    HashMap<Trap, LinkedList<Trap>> splitTraps = new HashMap<Trap, LinkedList<Trap>>();

    /**
     * Performs instrumentation for a single node.
     */
    public void BallLarusInstrument(CFGNode node, Local myPT)
    {
        //Collect outgoing edges for the node
        Collection<CFGEdge> edges = node.getSuccEdges().values();
        if (node.isDummy())
        {
            //No need to look for branches
            for (CFGEdge currEdge : node.getSuccEdges().values())
            {
                if (currEdge.getInstrVal() == 0) continue;
                if (currEdge.isDummyEdge()) continue;
                //
                instrumentEdge(currEdge, null);
            }
            return;
        }

        //Enumerate branches in the block
        List<Unit> branches = new LinkedList<Unit>();
        Iterator<Unit> unitIt = node.getBlock().iterator();
        while (unitIt.hasNext())
        {
            Unit currUnit = null;
            try
            {
                currUnit = unitIt.next(); 
                //TODO: we get a bug here sometimes. Caused by: java.util.NoSuchElementException: HashChain.LinkIterator.next() reached end of chain without reaching specified tail unit
                //Please report if this happens!
            }
            catch (NoSuchElementException e)
            {
                G.v().out.println("ERROR: Hit NoSuchElementException.  Please report this!");
                G.v().out.println("Printing block: " + node.getBlockID());
                break;
            }
            if (currUnit.branches())
            {
                branches.add(currUnit);
            }
        }

        //Match edges to branches to adjust appropriately
        for (CFGEdge currEdge : edges)
        {
            Trap t = null;
            if (currEdge.getInstrVal() == 0) continue; //No need to instrument
            if (currEdge.isDummyEdge()) continue; //No actual instrumentation
            if (currEdge.getTo().isDummy()) //Going to dummy exit block
            {
                instrumentEdge(currEdge, null);
                continue;
            }
            else if ((t = currEdge.getTrap()) != null) //Exceptional control-flow block
            {
                //Handle trap edge appropriately
                instrumentEdge(currEdge, t.getHandlerUnitBox());
                visitedTraps.add(t);
                continue;
            }
            else if(blockFallsThrough(currEdge.getFrom().getBlock(), currEdge.getTo().getBlock())) //Special case - don't bother adjusting jump
            {
                //Add instrumentation to end of from block
                instrumentEdge(currEdge, null);
                continue;
            }
            //Match branch to edge and adjust jump
            nextedge:
                for (Unit candUnit : branches)
                {
                    if (currEdge.getFrom().getBlock().getBody().getUnits().contains(candUnit))
                    {
                        //Found a match; instrument appropriately
                        List<UnitBox> targets = candUnit.getUnitBoxes();
                        for (UnitBox ub : targets)
                        {
                            if (isInBlock(currEdge.getTo().getBlock(), ub.getUnit()))
                            {
                                if(doDebug)
                                {
                                    G.v().out.println("Sanity check! " + isInBlock(currEdge.getTo().getBlock(), ub.getUnit()));
                                    G.v().out.println("currEdge: " + currEdge + " currUnit: ");
                                    printUnit(ub.getUnit());
                                }
                                instrumentEdge(currEdge, ub);
                                break;
                            }
                        }
                        break nextedge;
                    }
                }
        }

        //Do the same for our backedges
        for (CFGEdge backEdge : node.getBackedges())
        {
            nextedge:
                for (Unit candUnit : branches)
                {
                    if (backEdge.getFrom().getBlock().getBody().getUnits().contains(candUnit))
                    {
                        //Found a match; instrument appropriately
                        List<UnitBox> targets = candUnit.getUnitBoxes();
                        for (UnitBox ub : targets)
                        {
                            if (isInBlock(backEdge.getTo().getBlock(), ub.getUnit()))
                            {

                                if(doDebug)
                                {
                                    G.v().out.println("Sanity check! " + isInBlock(backEdge.getTo().getBlock(), ub.getUnit()));
                                    G.v().out.println("currEdge: " + backEdge + " currUnit: ");
                                    printUnit(ub.getUnit());
                                }
                                instrumentEdge(backEdge, ub);
                                break;
                            }
                        }
                        break nextedge;
                    }
                }
        }
    }

    /**
     * Returns true iff a block falls through to the next
     * @return true iff a block falls through to the next
     */
    public boolean blockFallsThrough(Block prior, Block latter)
    {
        Chain<Unit> units = prior.getBody().getUnits();
        if (units.getSuccOf(prior.getTail()) == latter.getHead()) return true;
        return false;
    }

    /**
     * Instruments an edge for us with the appropriate value.
     * Behavior changes depending on the type of edge and where
     * the original jump was to (if null, it's a dummy edge.)
     */
    public void instrumentEdge(CFGEdge edge, UnitBox originalJump)
    {
        Block succ = edge.getTo().getBlock();
        Block pred = edge.getFrom().getBlock();
        long addVal = edge.getInstrVal();
        InstanceFieldRef pathValRef = Jimple.v().newInstanceFieldRef(myPT, pathVal.makeRef());
        Chain<Unit> instrumentation = new HashChain<Unit>();

        if (!edge.isBackedge())
        {
            //Add instrumentation to increment our path value
            AddExpr pathAdd = Jimple.v().newAddExpr(pathCounter, LongConstant.v(addVal)); //Add increment to our counter
            AssignStmt pathAssign = Jimple.v().newAssignStmt(pathCounter, pathAdd); //Assign result back t
            instrumentation.add(pathAssign);
        }
        else
        {
            //Add instrumentation to reset our path value
            AssignStmt pathReset = Jimple.v().newAssignStmt(pathCounter, LongConstant.v(edge.getInstrVal())); //Reset path counter
            instrumentation.add(pathReset);
        }        

        //Instrumentation varies based on if node is dummy or if edge is backedge or falls through

        if (pred.getIndexInMethod() == ENTRY_DUMMY)
        {
            //Add instrumentation to beginning of to block
            //Skip parameters
            Unit blockHead = succ.getHead();
            Unit insertPoint = succ.getBody().getUnits().getSuccOf(blockHead);
            for(int i = 0; i < methodBody.getMethod().getParameterCount(); i++)
            {
                insertPoint = succ.getBody().getUnits().getSuccOf(insertPoint);
            }
            Chain<Unit> units = succ.getBody().getUnits();
            units.insertBefore(instrumentation, insertPoint);
        }
        else if (succ.getIndexInMethod() == EXIT_DUMMY)
        {
            //Add instrumentation to end of from block, just before return
            Unit blockTail = pred.getTail();
            Unit insertPoint = pred.getBody().getUnits().getPredOf(blockTail);
            Chain<Unit> units = pred.getBody().getUnits();
            units.insertAfter(instrumentation, insertPoint);
        }
        else if (originalJump == null)
        {
            //Block falls through; instrument at end of former
            Unit blockTail = pred.getTail();
            Chain<Unit> units = pred.getBody().getUnits();
            units.insertAfter(instrumentation, blockTail);
        }
        else
        {
            //Add jump to new instrumentation
            if (edge.getTrap() != null)
            {
                //We're instrumenting a trap edge
                if(doDebug) G.v().out.println("Instrumenting trapped edge " + edge.toString());
                Unit originalTarget = originalJump.getUnit();
                if (originalTarget != edge.getTrap().getHandlerUnit()) throw new RuntimeException("Trap target mismatch!");
                Unit target = originalTarget;

                //Walk past any nops to our caught-exception identity statement
                while (!(target instanceof IdentityStmt)) 
                {
                    target = succ.getBody().getUnits().getSuccOf(target);
                }
                IdentityStmt caughtException = (IdentityStmt) target;
                target = succ.getBody().getUnits().getSuccOf(target);
                GotoStmt jumpToOld = Jimple.v().newGotoStmt(target);

                //Add caught-exception identity in our instrumentation block
                soot.Local catchRefLocal = (Local) caughtException.getLeftOp();
                soot.jimple.CaughtExceptionRef caughtRef = soot.jimple.Jimple.v().newCaughtExceptionRef();
                IdentityStmt newCaughtException = Jimple.v().newIdentityStmt(catchRefLocal, caughtRef);

                instrumentation.addFirst(newCaughtException);
                instrumentation.addLast(jumpToOld);
                Chain<Unit> units = succ.getBody().getUnits();
                units.insertAfter(instrumentation, codeSinkStmt);

                //Update trap handler
                edge.getTrap().setHandlerUnit(newCaughtException);

                if(doDebug) G.v().out.println("Trap target: " + originalTarget + ", remaining nodes: " + trapMap.get(originalTarget));
                
                //Decrement counter for jumps to this target
                trapMap.put(originalTarget, (trapMap.get(originalTarget) - 1));

                //Remove original catch identity if nothing jumps there anymore
                if(trapMap.get(originalTarget) == 0)
                {
                    succ.remove(caughtException);
                }
                else
                {
                    if (doDebug)
                    {
                        G.v().out.println("Remaining traps: " + trapMap.get(originalTarget));
                        G.v().out.println("Remaining jumps to here: " + originalTarget.getBoxesPointingToThis());
                    }
                }
            }
            else
            {
                //Adjust original jump
                GotoStmt jumpToOld = Jimple.v().newGotoStmt(originalJump.getUnit());
                if (doDebug) G.v().out.println("Standard goto jump target: " + originalJump.getUnit());
                NopStmt dummyStmt = Jimple.v().newNopStmt();
                instrumentation.add(jumpToOld);
                instrumentation.addFirst(dummyStmt);
                Chain<Unit> units = succ.getBody().getUnits();
                units.insertAfter(instrumentation, codeSinkStmt);
                originalJump.setUnit(dummyStmt);
            }
        }
    }



    /**
     * Returns whether or not a given unit is in a block.
     * I'm not a fan of it at the moment though, since it just walks the block...
     */
    private boolean isInBlock(Block b, Unit u)
    {
        Iterator<Unit> myIt = b.iterator();
        if (myIt == null) return false;
        while (myIt.hasNext())
        {
            Unit curr = myIt.next();
            if (curr == u) return true;
        }
        return false;
    }

    /**
     * Auxilliary method to print a unit cleanly.
     */
    private void printUnit(Unit u)
    {
        NormalUnitPrinter up = new NormalUnitPrinter(methodBody);
        u.toString(up);
        if (doDebug) G.v().out.println(up.toString());        
    }

    /**
     * Mapping between basic blocks and their associated CFG nodes.
     */
    private HashMap<Block, CFGNode> blockNodes;

    /**
     * Inserts a block into our CFG.
     */
    private void buildCFG(Block b)
    {
        CFGNode blockNode;
        //Check if we've visited this block already
        if (blockNodes.containsKey(b))
        {
            blockNode = blockNodes.get(b);            
        }
        else
        {
            blockNode = new CFGNode(b);
        }
        if (b.getPreds().size() == 0)
        {
            //Entry block; link to dummy entry iff not a catch/finally block
            //Also check if dead code generated by Jimple
            //TODO: Confirm this is never going to miss real code.
            if (b.getIndexInMethod() == 0)
            {
                if (doDebug)
                {
                    G.v().out.println("Entry block:  " + b.getIndexInMethod());
                    G.v().out.println(b.toString());
                }
                makeCFGEdge(enterNode, blockNode);
            }
            else if (!isCatchOrFinallyBlock(b))
            {
                //Dead code, do not make block
                return;
            }
        }

        //Not a dead block; add node to list
        blockNodes.put(b, blockNode);

        //Link to successors
        List<Block> succs = b.getSuccs();
        if (succs.size() == 0)
        {
            //Exit node; link to dummy exit if throws and not caught
            //Or if does not throw
            boolean isTrapped = false;
            if (b.getTail() instanceof ThrowStmt)
            {
                for (Trap t : b.getBody().getTraps())
                {
                    if (isInTrap(b, t))
                    {
                        //All good; break
                        isTrapped = true;
                        break;
                    }
                }
                if (!isTrapped) makeCFGEdge(blockNode, exitNode);
            }
            else
            {
                makeCFGEdge(blockNode, exitNode);
            }
        }
        else
        {
            //Create edges between this block and its successors
            for (Block currBlock : succs)
            {
                CFGNode currNode;
                if (blockNodes.containsKey(currBlock))
                {
                    currNode = blockNodes.get(currBlock);
                }
                else
                {
                    currNode = new CFGNode(currBlock);
                    blockNodes.put(currBlock, currNode);
                }
                makeCFGEdge(blockNode, currNode);
            }
        }

    }

    /**
     * Returns true iff a block is the _start_ of a catch or finally block.
     * Not all that efficient.
     * @return Whether or not the block is the start of a catch or finally block.
     */
    private boolean isCatchOrFinallyBlock(Block b)
    {
        for (Trap t : b.getBody().getTraps())
        {
            if (t.getHandlerUnit() == b.getHead())
            {
                return true;
            }
        }
        return false;
    }

    /**
     * Generate an edge in our CFG between two nodes.
     */
    private CFGEdge makeCFGEdge(CFGNode from, CFGNode to)
    {
        CFGEdge newEdge = new CFGEdge(from, to);
        from.addSuccEdge(newEdge);
        from.addSucc(to);
        to.addPredEdge(newEdge);
        to.addPred(from);
        edges.add(newEdge);
        return newEdge;
    }

    /**
     * Performs our Ball-Larus instrumentation:
     * -Constructs CFG
     * -Calculates weights
     * -Calculates instrumentation values
     * -Instruments Jimple
     * -Prints results data to other files
     */
    public void BallLarus(List<Block> blocks, Local myPT)
    {
        //Set up start of CFG
        enterNode = new CFGNode(DUMMY_TYPE.ENTRY);
        exitNode = new CFGNode(DUMMY_TYPE.EXIT);
        blockNodes.put(enterNode.getBlock(), enterNode);
        blockNodes.put(exitNode.getBlock(), exitNode);

        //Construct CFG with info from each block
        for (int i = (blocks.size() - 1); i >=0; i--)
        {
            buildCFG(blocks.get(i));
        }

        //Add exception edges
        setExceptions();

        //Remove backedges
        fixBackedges();

        //Make weighted list of edges 
        ArrayList<CFGEdge> orderedList = new ArrayList<CFGEdge>();
        orderedList.addAll(edges);

        //Calculate spanning tree
        //HashSet<CFGEdge> spanningTree = buildSpanningTree(enterNode, exitNode, orderedList);
        HashSet<CFGEdge> spanningTree = buildWeightTree(orderedList);
        if (doDebug)
        {
            G.v().out.println("---------SPAN TREE--------");
            G.v().out.println(spanningTree);
            G.v().out.println("---------END  TREE--------");
        }


        //Calculate initial instrumentation values by walking backwards
        HashSet<CFGNode> visited = new HashSet<CFGNode>();
        Queue<CFGNode> topQueue = new LinkedList<CFGNode>();
        topQueue.add(exitNode);
        while (!topQueue.isEmpty())
        {
            CFGNode currNode = topQueue.poll();
            if (!visited.contains(currNode))
            {
                topQueue.addAll(currNode.getPreds());
                visited.add(currNode);
                BallLarus_stub(currNode);
            }
        }

        //Calculate final instrumentation values
        calculateInstruments(spanningTree);

        //Instrumentation pass
        for (CFGNode n : blockNodes.values())
        {
            BallLarusInstrument(n, myPT);
        }


        //Print out metadata to file
        SootMethod method = blocks.get(0).getBody().getMethod();
        String fileName = SootCSI.outDir + method.getDeclaringClass().getName() + "_pt";
        try
        {
            PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter( 
                    fileName, true)));
            out.println("#");
            out.println(method.getSignature());

            //Print blocks
            for (Block b : blocks)
            {
                printBlockLineData(out, b);
            }
            out.println("$");

            //Print edges
            orderedList.addAll(backEdges);
            printEdgeData(out, orderedList);
            out.flush();
            out.close();
        }
        catch (IOException e)
        {
            if (doDebug) G.v().out.println("Warning: Could not open file " + fileName);
        }

    }


    /**
     * Auxilliary method to print out information about each block's line numbers.
     */
    private void printBlockLineData(PrintWriter out, Block b)
    {
        out.print(b.getIndexInMethod());
        for (Unit u = b.getHead() ; u != b.getSuccOf(b.getTail()) ; u = b.getSuccOf(u))
        {
            SourceLnPosTag lineInfoTag = null;
            Iterator<Tag> startTagIterator = u.getTags().iterator();
            while (startTagIterator.hasNext())
            {
                Tag tempTag = startTagIterator.next();
                if (tempTag instanceof SourceLnPosTag)
                {
                    lineInfoTag = (SourceLnPosTag) tempTag;
                    break;
                }
            }

            if (lineInfoTag == null) continue; //No source info!
            out.print("|");
            out.print(lineInfoTag.startLn());
        }
        out.println();
    }

    /**
     * Auxiliary method to print out info on each edge and its 
     * instrumentation value.
     */
    private void printEdgeData(PrintWriter out, List<CFGEdge> edges)
    {
        for (CFGEdge e : edges)
        {
            if (e.isDummyEdge()) continue;
            if (e.isBackedge())
            {
                out.print(e.getFrom().getBlockID() + "~>" + e.getTo().getBlockID());
                out.print("|");
                out.print(e.getAssociatedDummyEdge().getInstrVal() + "$" + e.getAssociatedDummyEdge().getInstrVal());
            }
            else
            {
                out.print(e.getFrom().getBlockID() + "->" + e.getTo().getBlockID());
                out.print("|");
                out.print(e.getVal() + "$" + e.getInstrVal());
            }
            out.println();
        }
    }

    private Local pathCounter, ptUse, myPT;
    private Body methodBody;
    private ZonedBlockGraph blockGraph;
    private Unit codeSinkStmt;
    private boolean doDebug = false;

    /**
     * Checks for traps that catch the first unit and adjusts them
     * to begin after our initialization instrumentation.
     * 
     * Also inserts assignments to get our PathTrace object after
     * the caught-exception identity of a handler.
     */
    private void fixTraps(Body methodBody, Unit safeStart)
    {
        for (Trap t : methodBody.getTraps())
        {
            Unit start = t.getBeginUnit();
            if (start == methodBody.getUnits().getFirst())
            {
                //Move to start after init
                t.setBeginUnit(safeStart);
            }
            Unit handler = t.getHandlerUnit();
            
            //Insert after exception identity
            StaticFieldRef useRef = Jimple.v().newStaticFieldRef(inUse.makeRef());
            AssignStmt tempAssign = Jimple.v().newAssignStmt(ptUse, useRef);
            methodBody.getUnits().insertAfter(tempAssign, handler);
        }
    }

    /**
     * Main method for our instrumentation.
     */
    public PathAnalysis(ZonedBlockGraph expandedBlockGraph) {
        if (instruments == null) instruments = new HashMap<CFGNode, Integer>();
        blockNodes = new HashMap<Block, CFGNode>();
        methodBody = expandedBlockGraph.getBody();
        if(methodBody.getMethod().getDeclaringClass().getName().contains("uwmadison.soot.analysis")) return;
        SootMethod myMethod = methodBody.getMethod();
        blockGraph = new ZonedBlockGraph(methodBody);
        Chain<Unit> units = blockGraph.getBody().getUnits();

        Stmt blockHead = (Stmt) units.getFirst();      
        codeSinkStmt = units.getLast();

        //Set up our PathTrace object + temporary long
        Chain<Local> locals = methodBody.getLocals();
        String ptName = myMethod.getName() + "__PT_holder";
        Type ptType = pathClass.getType();
        myPT = Jimple.v().newLocal(ptName, ptType);

        pathCounter = Jimple.v().newLocal("_BL_path_counter", LongType.v());

        ptUse = Jimple.v().newLocal("_BL_inuse_temp", BooleanType.v());

        //Add them to the method's locals
        locals.add(myPT);
        locals.add(pathCounter);
        locals.add(ptUse);

        //Call to factory method
        StaticInvokeExpr constructExpr = Jimple.v().newStaticInvokeExpr(
                getRef.makeRef(), StringConstant.v(myMethod.getSignature()));

        //Call to PathTrace stack-freeing method
        InvokeExpr closeExpr = Jimple.v().newStaticInvokeExpr(freeStack.makeRef());

        //Initial null-assignment for path trace object
        AssignStmt nullInit = Jimple.v().newAssignStmt(myPT, NullConstant.v());

        //Initial zero-assignment for path val and use bool
        AssignStmt pathInit = Jimple.v().newAssignStmt(pathCounter, LongConstant.v(0));
        AssignStmt ptInit = Jimple.v().newAssignStmt(ptUse, IntConstant.v(0));

        //Don't bother instrumenting if only one basic block
        if (blockGraph.getBlocks().size() > 1) 
        {
            //Perform instrumentation
            BallLarus(blockGraph.getBlocks(), myPT);
        }


        //Refresh the block graph with our instrumented Jimple code
        blockGraph = new ZonedBlockGraph(methodBody);
        blockHead = (Stmt) units.getFirst();
        Stmt insertPoint = blockHead;

        //We need to take the type of the method into account when
        //determining where to place our instrumentation variables
        if (methodBody.getMethod().getParameterCount() > 0)
        {
            //Check if static (eg. no @this-assignment)
            int offset = 0;
            if (methodBody.getMethod().isStatic()) offset = 1;
            //Walk past parameters
            for (int i = offset; i < methodBody.getMethod().getParameterCount(); i++)
            {
                insertPoint = (Stmt) units.getSuccOf(insertPoint);
            }
        }
        else
        {
            //Can just start directly at method entry
            insertPoint = blockHead;
        }
        
        //Boundary point for assigning traps
        NopStmt trapEdge = Jimple.v().newNopStmt();

        //Insert initializations    
        units.insertAfter(trapEdge, insertPoint); 
        units.insertAfter(nullInit, insertPoint);     
        units.insertAfter(pathInit, insertPoint);     
        units.insertAfter(ptInit, insertPoint);       

        //For determining trap range
        Unit startTrap = units.getSuccOf(trapEdge);

        //Instrument catch blocks to check PT status
        //And handle case of traps starting with first unit
        fixTraps(methodBody, startTrap);

        //Chain of instructions to handle uncaught exception case
        Chain<Unit> uncaughtChain = new HashChain<Unit>();        

        //Catch uncaught exceptions and store result
        Local catchRefLocal = soot.jimple.Jimple.v().newLocal("e", soot.RefType.v("java.lang.Throwable"));
        methodBody.getLocals().add(catchRefLocal);
        CaughtExceptionRef caughtRef = soot.jimple.Jimple.v().newCaughtExceptionRef();
        Stmt caughtIdentity = soot.jimple.Jimple.v().newIdentityStmt(catchRefLocal, caughtRef);
        uncaughtChain.add(caughtIdentity);

        //Insert call to factory method for PathTrace object
        uncaughtChain.add(Jimple.v().newAssignStmt(myPT, constructExpr));

        //Assign current path val to PathTrace object
        InstanceFieldRef pathValRef = Jimple.v().newInstanceFieldRef(myPT, pathVal.makeRef());
        AssignStmt linkStmt = Jimple.v().newAssignStmt(pathValRef, pathCounter);
        uncaughtChain.add(linkStmt);

        //Check if we're the main method
        if (myMethod.isMain())
        {
            //Special case: dump info
            StaticInvokeExpr dumpExpr = Jimple.v().newStaticInvokeExpr(
                    dumpAllTraces.makeRef());
            Stmt dumpStmt = Jimple.v().newInvokeStmt(dumpExpr);
            uncaughtChain.add(dumpStmt);
        }

        //Throw exception
        ThrowStmt throwStmt = Jimple.v().newThrowStmt(catchRefLocal);
        uncaughtChain.add(throwStmt);

        //Insert instrumentation        
        Unit endStmt = units.getLast();
        units.insertAfter(uncaughtChain, endStmt);
        
        //Move "endStmt" one forward because newTrap is exclusive on the end of the trap
        endStmt = units.getSuccOf(endStmt);

        //Add trap to instrumentation
        methodBody.getTraps().add(Jimple.v().newTrap(throwableClass, startTrap, endStmt, caughtIdentity));

        //Finally, insert calls to PathTrace freeStack() before returns
        Iterator<Unit> stmtIt = units.snapshotIterator();
        while (stmtIt.hasNext())
        {
            Unit u = (Unit) stmtIt.next();
            if ((u instanceof ReturnStmt) || (u instanceof RetStmt)  || (u instanceof ReturnVoidStmt))
            {
                //Quickly check if we have anything in stack that needs freeing
                IfStmt useCheck = Jimple.v().newIfStmt(Jimple.v().newEqExpr(ptUse,IntConstant.v(0)), u);
                units.insertBefore(useCheck, u);
                
                //Insert close op just before
                InvokeStmt closeInvoke = Jimple.v().newInvokeStmt(closeExpr);
                units.insertBefore(closeInvoke, u);
                
                //Ensure proper jump target
                useCheck.setTarget(u);
            }
        }

        if(doDebug)
        {
            G.v().out.println("--------TRANSFORMED--------");
            for (Unit u : expandedBlockGraph.getBody().getUnits())
            {
                printUnit(u);
            }
            G.v().out.println("----END TRANSFORMED--------");
        }
    }
}